import pandas as pd

PARA_BUY_THRESHOLD = 0.003
PARA_SELL_THRESHOLD = 0
PARA_SHORT_THRESHOLD = -0.003
PARA_COVER_THRESHOLD = 0
COMMISSION_RATE = 0.00015   # 单边手续费率


def make_strategy_result(pred_data):
    secs = pred_data["sec_id"].drop_duplicates().to_list()
    pred_data["date"] = pred_data["datetime"].dt.date

    for sec_id in secs:
        print(sec_id)
        sec_data = pred_data.loc[pred_data["sec_id"] == sec_id]
        data_profits = {}

        current_symbol = None
        trades = []
        last_date = None
        date_commission = 0
        date_profit = 0
        last_pos_price = 0
        last_price = 0
        pos = 0
        date_pnls = []
        for n, rows in sec_data.iterrows():
            pred_y = rows["pred_y"]
            date = rows["date"]
            symbol = rows["symbol"]
            price = rows["close"]
            dt = rows["datetime"]
            if not last_date:
                last_date = date
            if not current_symbol:
                current_symbol = symbol
            elif symbol != current_symbol:
                # 合约切换，如果有持仓，则按原价平掉
                if pos == 1:
                    date_commission += COMMISSION_RATE
                    date_profit += (last_price - last_pos_price) / last_pos_price - COMMISSION_RATE

                    trade = {"direction": "short", "offset": "close", "price": last_price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
                elif pos == -1:
                    date_commission += COMMISSION_RATE
                    date_profit += (last_pos_price - last_price) / last_pos_price - COMMISSION_RATE
                    trade = {"direction": "long", "offset": "close", "price": last_price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
                last_pos_price = 0
                pos = 0
                current_symbol = symbol
            if pos == 1:
                if pred_y < PARA_SELL_THRESHOLD:
                    # 平多
                    date_commission += COMMISSION_RATE
                    date_profit += (price - last_pos_price) / last_pos_price - COMMISSION_RATE
                    last_pos_price = 0
                    pos = 0
                    trade = {"direction": "short", "offset": "close", "price": price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
                if pred_y < PARA_SHORT_THRESHOLD:
                    pos = -1
                    date_commission += COMMISSION_RATE
                    last_pos_price = price
                    trade = {"direction": "short", "offset": "open", "price": price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
            elif pos == -1:
                if pred_y > PARA_COVER_THRESHOLD:
                    # 平空
                    date_commission += COMMISSION_RATE
                    date_profit += (last_pos_price - price) / last_pos_price - COMMISSION_RATE
                    last_pos_price = 0
                    pos = 0
                    trade = {"direction": "long", "offset": "close", "price": price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
                if pred_y > PARA_BUY_THRESHOLD:
                    pos = 1
                    date_commission += COMMISSION_RATE
                    last_pos_price = price
                    trade = {"direction": "long", "offset": "open", "price": price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
            else:
                if pred_y > PARA_BUY_THRESHOLD:
                    pos = 1
                    date_commission += COMMISSION_RATE
                    last_pos_price = price
                    trade = {"direction": "long", "offset": "open", "price": price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
                elif pred_y < PARA_SHORT_THRESHOLD:
                    pos = -1
                    date_commission += COMMISSION_RATE
                    last_pos_price = price
                    trade = {"direction": "short", "offset": "open", "price": price, "symbol": symbol, "dt": dt}
                    trades.append(trade)
            # 换日
            if last_date and date != last_date:
                if pos == 1:
                    pos_profit = (price - last_pos_price) / last_pos_price
                elif pos == -1:
                    pos_profit = (last_pos_price - price) / last_pos_price
                else:
                    pos_profit = 0
                date_profit += pos_profit
                date_pnls.append({"date": date, "profit": date_profit, "commission": date_commission, "symbol": symbol})
                date_commission = 0
                date_profit = 0
                if pos:
                    last_pos_price = price
                else:
                    last_pos_price = 0
                last_date = date
            last_price = price

        df = pd.DataFrame(date_pnls)
        df["total_profit"] = df["profit"].cumsum()
        df["total_commission"] = df["commission"].cumsum()
        #df.to_csv(f"D:\daily work\ml\\xgb\\strategy_{sec_id}.csv", index=False)

        trade_df = pd.DataFrame(trades)
        #trade_df.to_csv(f"D:\daily work\ml\\xgb\\trades_{sec_id}.csv", index=False)
        return df, trade_df


if __name__ == "__main__":
    pred_data = pd.read_csv("D:\daily work\ml\\xgb\\result.csv")
    pred_data["datetime"] = pd.to_datetime(pred_data["datetime"])
    df, trade_df = make_strategy_result(pred_data)
